use super::*;
use priority_queue::PriorityQueue;
use std::{sync::RwLock, time::Duration};
use std::any::Any;
use tokio::sync::mpsc;
use crate::message::{Message, DynamicMessage};
use crate::errors::SendError;
use super::{MailboxBehavior, MailboxConfig, MailboxStatus};
use std::collections::BinaryHeap;
use std::cmp::Ordering;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum MessagePriority {
    System = 0,
    High = 1,
    Normal = 2,
    Low = 3,
    Background = 4,
}

impl Default for MessagePriority {
    fn default() -> Self {
        MessagePriority::Normal
    }
}

pub trait PrioritizedMessage {
    fn priority(&self) -> MessagePriority;
}

impl PrioritizedMessage for Message {
    fn priority(&self) -> MessagePriority {
        match self {
            Message::System(_) => MessagePriority::System,
            Message::User(msg) => msg.priority(),
        }
    }
}

#[derive(Clone)]
struct PriorityMessage {
    priority: u8,
    message: DynamicMessage,
}

impl PartialEq for PriorityMessage {
    fn eq(&self, other: &Self) -> bool {
        self.priority == other.priority
    }
}

impl Eq for PriorityMessage {}

impl PartialOrd for PriorityMessage {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for PriorityMessage {
    fn cmp(&self, other: &Self) -> Ordering {
        self.priority.cmp(&other.priority)
    }
}

pub struct PriorityMailbox {
    sender: mpsc::UnboundedSender<PriorityMessage>,
    receiver: mpsc::UnboundedReceiver<PriorityMessage>,
    config: MailboxConfig,
    stats: MailboxStatus,
    queue: BinaryHeap<PriorityMessage>,
}

impl PriorityMailbox {
    pub fn new(config: MailboxConfig) -> Self {
        let (sender, receiver) = mpsc::unbounded_channel();
        Self {
            sender,
            receiver,
            config,
            stats: MailboxStatus {
                messages_received: 0,
                messages_processed: 0,
            },
            queue: BinaryHeap::new(),
        }
    }
}

#[async_trait::async_trait]
impl MailboxBehavior for PriorityMailbox {
    async fn post<T: Any + Send + Clone + 'static>(&self, msg: Message<T>) -> Result<(), SendError> {
        let dynamic_msg = msg.into_dynamic();
        self.post_dynamic(dynamic_msg).await
    }

    async fn post_dynamic(&self, msg: DynamicMessage) -> Result<(), SendError> {
        let priority_msg = PriorityMessage {
            priority: msg.priority,
            message: msg,
        };
        self.sender.send(priority_msg).map_err(|_| SendError::MailboxFull)
    }

    async fn receive(&self) -> Result<Option<DynamicMessage>, SendError> {
        if let Some(msg) = self.queue.pop() {
            Ok(Some(msg.message))
        } else if let Some(msg) = self.receiver.recv().await {
            Ok(Some(msg.message))
        } else {
            Ok(None)
        }
    }

    fn stats(&self) -> MailboxStatus {
        self.stats.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mailbox::MailboxConfig;
    use crate::message::Message;
    use std::sync::Arc;

    #[tokio::test]
    async fn test_priority_mailbox_send_and_receive() {
        let config = MailboxConfig::default();
        let mailbox = PriorityMailbox::new(config);
        let msg = Message::new("Test message");

        // Send a message
        assert!(mailbox.send(msg.clone()).await.is_ok());

        // Receive the message
        let received_msg = mailbox.receive().await.unwrap();
        assert_eq!(received_msg, Some(msg));
    }

    #[tokio::test]
    async fn test_priority_mailbox_receive_empty() {
        let config = MailboxConfig::default();
        let mailbox = PriorityMailbox::new(config);

        // Try to receive a message when none have been sent
        let received_msg = mailbox.receive().await.unwrap();
        assert_eq!(received_msg, None);
    }
} 